import os
import cv2
from PIL import Image
import sys
import numpy as np
import random
sys.path.append('..')
from config import *
import torch

from torch.utils.data import Dataset
import torchvision.transforms as transform



class PlateDataset(Dataset):
    """
    txtfile存在config.py文件中
    """
    def __init__(self, txtfile) -> None:    # image_12_47.txt or image_24_94.txt
        super(PlateDataset, self).__init__()

        self.trans = transform.Compose([
            transform.ToTensor()
        ])

        self.file = os.path.dirname(os.getcwd()) + '/Data_preprocess/'\
                    + str(txtfile).split('.')[0].split('_', 1)[1] + '/' + txtfile   # 文件路径

        # self.data = []

        with open(self.file, 'r') as f:
            self.img_idx = [x.strip().split(' ')[0] for x in f.readlines()]  # 图片输入数据

        with open(self.file, 'r') as f:
            self.label = np.array([int(x.strip().split(' ')[1]) for x in f.readlines()]) # 分类标签

        self.reg = []
        with open(self.file, 'r') as f:
            for x in f.readlines():
                if x.strip().split(' ')[1] != '0':
                    self.reg.append(np.array(list(map(float, x.strip().split(' ')[2:])))) # 回归参数
                else:
                    self.reg.append(np.zeros((4, )))
                    
        # 随机打乱
        random.seed(1)
        random.shuffle(self.img_idx)
        random.seed(1)
        random.shuffle(self.label)
        random.seed(1)
        random.shuffle(self.reg)
        
        
        # for img, lb, reg in zip(self.img, self.label, self.reg):
        #     data_ = dict()
        #     data_['img'] = img
        #     data_['label'] = lb
        #     data_['reg'] = reg
        #     self.data.append(data_)
        # random.shuffle(self.data)   # 将数据存储成字典数组的形式并进行打乱

    def __len__(self) -> int:
        return len(self.label)

    def __getitem__(self, index: int):
        global_path = os.path.dirname(os.getcwd()) + '/Data_preprocess' + self.img_idx[index].strip('.')
        self.img = cv2.imread(f'{global_path}') # int8
        # self.img = Image.open(f'{global_path}')
        self.img = self.trans(self.img)
        # print(self.img.shape)

        # self.label = torch.from_numpy(self.label)
        self.reg[index] = torch.from_numpy(self.reg[index])
        return self.img, self.label[index], self.reg[index]


if __name__ == '__main__':
    pass